#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
创建单个conda环境的脚本，使用清华镜像源
"""

import os
import sys
import subprocess
import argparse
from pathlib import Path

# 清华镜像源URL
TSINGHUA_MIRROR = "https://mirrors.tuna.tsinghua.edu.cn/anaconda"

def list_available_environments():
    """列出可用的环境文件"""
    env_dir = Path("conda_environments")
    if not env_dir.exists():
        print(f"错误: 找不到环境目录 {env_dir}")
        return []
    
    env_files = list(env_dir.glob("*.yml"))
    return [f.name for f in env_files]

def create_environment(env_file, env_name=None, use_tsinghua=True):
    """创建conda环境"""
    env_path = Path("conda_environments") / env_file
    
    if not env_path.exists():
        print(f"错误: 环境文件 {env_file} 不存在")
        return False
    
    print(f"正在使用文件 {env_file} 创建conda环境...")
    
    # 设置创建环境命令
    command = ["conda", "env", "create", "-f", str(env_path)]
    
    # 如果指定了环境名称，添加到命令中
    if env_name:
        command.extend(["-n", env_name])
        print(f"将使用自定义环境名称: {env_name}")
    
    # 添加清华镜像源参数
    if use_tsinghua:
        # 需要在环境变量中设置镜像源
        env = os.environ.copy()
        env["CONDA_PKGS_DIRS"] = os.path.expanduser("~/.conda/pkgs")
        
        # 设置conda配置使用清华镜像源
        subprocess.run(["conda", "config", "--set", "show_channel_urls", "yes"], 
                       check=False)
        subprocess.run(["conda", "config", "--add", "channels", f"{TSINGHUA_MIRROR}/cloud/conda-forge"], 
                       check=False)
        subprocess.run(["conda", "config", "--add", "channels", f"{TSINGHUA_MIRROR}/pkgs/free"], 
                       check=False)
        subprocess.run(["conda", "config", "--add", "channels", f"{TSINGHUA_MIRROR}/pkgs/main"], 
                       check=False)
        
        print("已配置清华镜像源")
        
        # 执行创建环境命令
        process = subprocess.run(command, env=env)
    else:
        process = subprocess.run(command)
    
    if process.returncode != 0:
        print(f"环境创建失败，返回码: {process.returncode}")
        return False
    
    print("环境创建成功!")
    return True

def main():
    parser = argparse.ArgumentParser(description="创建单个conda环境，使用清华镜像源")
    parser.add_argument("--env-file", type=str, help="环境文件名 (位于conda_environments目录下)")
    parser.add_argument("--env-name", type=str, help="指定创建的conda环境名称")
    parser.add_argument("--list", action="store_true", help="列出可用的环境文件")
    parser.add_argument("--no-tsinghua", action="store_true", help="不使用清华镜像源")
    
    args = parser.parse_args()
    
    if args.list:
        print("可用的环境文件:")
        for env in list_available_environments():
            print(f"- {env}")
        return
    
    if not args.env_file:
        # 如果没有指定环境文件，列出可用选项并让用户选择
        available_envs = list_available_environments()
        if not available_envs:
            return
        
        print("请选择要创建的环境文件:")
        for i, env in enumerate(available_envs):
            print(f"{i+1}. {env}")
        
        try:
            choice = int(input("请输入选项编号: "))
            if 1 <= choice <= len(available_envs):
                env_file = available_envs[choice-1]
            else:
                print("无效的选择")
                return
        except ValueError:
            print("请输入有效的数字")
            return
    else:
        env_file = args.env_file
    
    # 如果没有指定环境名称，询问用户是否要指定
    env_name = args.env_name
    if not env_name and input("是否要指定自定义环境名称? (y/n): ").lower() == 'y':
        env_name = input("请输入环境名称: ").strip()
        if not env_name:
            print("未指定环境名称，将使用配置文件中的默认名称")
    
    create_environment(env_file, env_name, not args.no_tsinghua)

if __name__ == "__main__":
    main() 